package cn.ocars.playgame.pencilstudy.model

import android.content.Context
import android.graphics.Bitmap
import android.util.Log
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import org.json.JSONObject
import java.io.IOException
import java.nio.ByteBuffer

/**
 * MiniCPM-V模型加载器
 * 负责加载和管理MiniCPM-V模型，提供模型推理接口
 */
class MiniCPMVModelLoader(private val context: Context) {
    private val TAG = "MiniCPMVModelLoader"
    
    private var miniCPMEngine: MiniCPMEngine? = null
    private var modelConfig: ModelConfig? = null
    private var isModelLoaded = false
    
    /**
     * 模型配置数据类
     */
    data class ModelConfig(
        val modelName: String,
        val modelVersion: String,
        val modelPath: String,
        val modelType: String,
        val modelDescription: String,
        val modelSize: String,
        val modelQuantization: String,
        val maxSequenceLength: Int,
        val temperature: Float,
        val topP: Float,
        val topK: Int,
        val repetitionPenalty: Float,
        val inputSize: Pair<Int, Int>,
        val strokeThreshold: Float,
        val confidenceThreshold: Float,
        val maxStrokes: Int,
        val minStrokes: Int,
        val structureWeight: Float,
        val proportionWeight: Float,
        val strokeOrderWeight: Float,
        val maxScore: Float,
        val minScore: Float
    )
    
    /**
     * 初始化模型
     */
    suspend fun initialize(): Boolean = withContext(Dispatchers.IO) {
        try {
            Log.d(TAG, "开始初始化MiniCPM-V模型...")
            
            // 加载模型配置
            loadModelConfig()
            
            // 创建MiniCPM引擎
            miniCPMEngine = MiniCPMEngine(context)
            
            // 初始化引擎
            val modelPath = modelConfig?.modelPath ?: "minicpm-v-2.5"
            val modelQuantization = modelConfig?.modelQuantization ?: "q4f16_0"
            val success = miniCPMEngine?.initialize(modelPath, modelQuantization) ?: false
            
            if (success) {
                isModelLoaded = true
                Log.d(TAG, "MiniCPM-V模型初始化成功")
            } else {
                Log.e(TAG, "MiniCPM-V模型初始化失败")
                isModelLoaded = false
            }
            
            success
        } catch (e: Exception) {
            Log.e(TAG, "MiniCPM-V模型初始化失败", e)
            isModelLoaded = false
            false
        }
    }
    
    /**
     * 加载模型配置
     */
    private fun loadModelConfig() {
        try {
            val configJson = context.assets.open("minicpm-v-config.json")
                .bufferedReader().use { it.readText() }
            
            val jsonObject = JSONObject(configJson)
            
            modelConfig = ModelConfig(
                modelName = jsonObject.getString("model_name"),
                modelVersion = jsonObject.getString("model_version"),
                modelPath = jsonObject.getString("model_path"),
                modelType = jsonObject.getString("model_type"),
                modelDescription = jsonObject.getString("model_description"),
                modelSize = jsonObject.getString("model_size"),
                modelQuantization = jsonObject.getString("model_quantization"),
                maxSequenceLength = jsonObject.getJSONObject("model_parameters").getInt("max_sequence_length"),
                temperature = jsonObject.getJSONObject("model_parameters").getDouble("temperature").toFloat(),
                topP = jsonObject.getJSONObject("model_parameters").getDouble("top_p").toFloat(),
                topK = jsonObject.getJSONObject("model_parameters").getInt("top_k"),
                repetitionPenalty = jsonObject.getJSONObject("model_parameters").getDouble("repetition_penalty").toFloat(),
                inputSize = Pair(
                    jsonObject.getJSONObject("image_processing").getJSONArray("input_size").getInt(0),
                    jsonObject.getJSONObject("image_processing").getJSONArray("input_size").getInt(1)
                ),
                strokeThreshold = jsonObject.getJSONObject("handwriting_recognition").getDouble("stroke_threshold").toFloat(),
                confidenceThreshold = jsonObject.getJSONObject("handwriting_recognition").getDouble("confidence_threshold").toFloat(),
                maxStrokes = jsonObject.getJSONObject("handwriting_recognition").getInt("max_strokes"),
                minStrokes = jsonObject.getJSONObject("handwriting_recognition").getInt("min_strokes"),
                structureWeight = jsonObject.getJSONObject("scoring_parameters").getDouble("structure_weight").toFloat(),
                proportionWeight = jsonObject.getJSONObject("scoring_parameters").getDouble("proportion_weight").toFloat(),
                strokeOrderWeight = jsonObject.getJSONObject("scoring_parameters").getDouble("stroke_order_weight").toFloat(),
                maxScore = jsonObject.getJSONObject("scoring_parameters").getDouble("max_score").toFloat(),
                minScore = jsonObject.getJSONObject("scoring_parameters").getDouble("min_score").toFloat()
            )
            
            Log.d(TAG, "模型配置加载成功: ${modelConfig?.modelName} ${modelConfig?.modelVersion}")
        } catch (e: IOException) {
            Log.e(TAG, "加载模型配置失败", e)
            // 使用默认配置
            modelConfig = ModelConfig(
                modelName = "MiniCPM-V",
                modelVersion = "2.5",
                modelPath = "minicpm-v-2.5",
                modelType = "multimodal",
                modelDescription = "MiniCPM-V 2.5是一个多模态大模型，支持图像和文本输入，适用于手写识别和评分任务",
                modelSize = "1.8B",
                modelQuantization = "q4f16_0",
                maxSequenceLength = 2048,
                temperature = 0.7f,
                topP = 0.9f,
                topK = 50,
                repetitionPenalty = 1.1f,
                inputSize = Pair(224, 224),
                strokeThreshold = 0.5f,
                confidenceThreshold = 0.7f,
                maxStrokes = 50,
                minStrokes = 1,
                structureWeight = 0.4f,
                proportionWeight = 0.3f,
                strokeOrderWeight = 0.3f,
                maxScore = 1.0f,
                minScore = 0.0f
            )
        }
    }
    
    /**
     * 释放模型资源
     */
    fun release() {
        if (isModelLoaded) {
            miniCPMEngine?.release()
            miniCPMEngine = null
            isModelLoaded = false
            Log.d(TAG, "MiniCPM-V模型资源已释放")
        }
    }
    
    /**
     * 检查模型是否已加载
     */
    fun isModelLoaded(): Boolean {
        return isModelLoaded
    }
    
    /**
     * 获取模型配置
     */
    fun getModelConfig(): ModelConfig? {
        return modelConfig
    }
    
    /**
     * 生成评分
     * @param prompt 提示词
     * @param image 图像数据
     * @return 模型生成的响应
     */
    suspend fun generateScore(prompt: String, image: ByteBuffer): String? = withContext(Dispatchers.IO) {
        if (!isModelLoaded) {
            Log.w(TAG, "模型未加载，无法生成评分")
            return@withContext null
        }
        
        try {
            Log.d(TAG, "正在使用MiniCPM-V模型生成评分...")
            
            // 使用MiniCPM引擎生成响应
            val response = miniCPMEngine?.generateScore(prompt, image)
            
            Log.d(TAG, "评分生成成功")
            response
        } catch (e: Exception) {
            Log.e(TAG, "生成评分失败", e)
            null
        }
    }
    
    /**
     * 生成评分（使用Bitmap）
     * @param prompt 提示词
     * @param bitmap 图像位图
     * @return 模型生成的响应
     */
    suspend fun generateScore(prompt: String, bitmap: Bitmap): String? = withContext(Dispatchers.IO) {
        if (!isModelLoaded) {
            Log.w(TAG, "模型未加载，无法生成评分")
            return@withContext null
        }
        
        try {
            Log.d(TAG, "正在使用MiniCPM-V模型生成评分...")
            
            // 使用MiniCPM引擎生成响应
            val response = miniCPMEngine?.generateScore(prompt, bitmap)
            
            Log.d(TAG, "评分生成成功")
            response
        } catch (e: Exception) {
            Log.e(TAG, "生成评分失败", e)
            null
        }
    }
}